# utils/memory_system.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
记忆系统
"""
import json
import logging
import os
import re
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

# 本地模块导入
try:
    from utils.llm_service import LLMService
except ImportError:
    LLMService = None  # type: ignore

from models import User, UserBuffer, Conversation, ConversationRecord, db
from utils import llm_service

logger = logging.getLogger("AI-MindCare-System-Memory")


class MemorySystem:
    """记忆系统，管理用户信息、咨询记录、对话历史等"""
    def __init__(self, llm_service: Optional[LLMService] = None):
        self.llm_service = llm_service
    
    def _now_utc(self) -> datetime:
        return datetime.now(timezone.utc)
    

  

    
    def _load_long_term_memory(self, user_id: uuid.UUID) -> Optional[Dict[str, Any]]:
        """加载用户的长期记忆"""
        user = User.query.get(user_id)
        if user and user.long_term_memory:
            return user.long_term_memory
        return None
    def _update_long_term_memory(self, user_id: uuid.UUID, long_term_memory: Dict[str, Any]):
        """更新用户的长期记忆"""
        user = User.query.get(user_id)
        if user:
            user.long_term_memory = long_term_memory
        else:
            new_user = User(id=user_id, long_term_memory=long_term_memory)
            db.session.add(new_user)
        try:
            db.session.commit()
            logger.info(f"用户 {user_id} 的长期记忆已更新到数据库")
        except Exception as e:
            db.session.rollback()
            logger.error(f"更新用户 {user_id} 长期记忆到数据库时出错: {e}")

    
    def _get_user_buffers(self) -> Dict[str, Dict[str, List[Any]]]:
        """从 UserBuffer 表获取所有用户的会话缓冲区"""
        buffers = {}
        records = UserBuffer.query.order_by(
            UserBuffer.conversation_id,
            UserBuffer.timestamp.desc()
        ).all()
        for record in records:
            user_id_str = str(record.user_id)
            if user_id_str not in buffers:
                buffers[user_id_str] = {}
            if record.conversation_id not in buffers[user_id_str]:
                buffers[user_id_str][record.conversation_id] = []
            buffers[user_id_str][record.conversation_id].append(record.message)
        return buffers


    
    def add_to_user_buffer(self, user_id: uuid.UUID, conversation_id: str, message: str, role: str):
        """向用户的会话缓冲区添加新消息"""
        new_conversation_record = ConversationRecord(
            user_id=user_id,
            conversation_id=conversation_id,
            role=role,
            content=message,
            message_id=str(uuid.uuid4()),
            created_at=datetime.now(timezone.utc)
        )
        db.session.add(new_conversation_record)
        new_buffer = UserBuffer(user_id=user_id,
            conversation_id=conversation_id,
            role=role,
            content=message,
            message_id=str(uuid.uuid4()),
            created_at=datetime.now(timezone.utc))
        db.session.add(new_buffer)

        try:
            db.session.commit()
            # 清理 UserBuffer 表中旧记录，只保留每个会话最近的 20 条
            old_records = UserBuffer.query.filter_by(
                user_id=user_id,
                conversation_id=conversation_id
            ).order_by(UserBuffer.timestamp.asc()).limit(
                db.session.query(UserBuffer).filter_by(
                    user_id=user_id,
                    conversation_id=conversation_id
                ).count() - 20
            )
            for record in old_records:
                db.session.delete(record)
            db.session.commit()
            logger.info(f"已为用户 {user_id} 会话 {conversation_id} 添加新消息")
        except Exception as e:
            db.session.rollback()
            logger.error(f"为用户 {user_id} 会话 {conversation_id} 添加新消息时出错: {e}")


    def get_user_buffer(self, conversation_id: str) -> List[str]:
        """获取用户某个会话的缓冲区"""

        records = UserBuffer.query.filter_by(
            conversation_id=conversation_id
        ).order_by(UserBuffer.created_at.asc()).all()

        return [{
            "role": record.role,
            "content": record.content,
            "message_id": record.message_id,
            "parent_message_id": record.parent_message_id,
            # "created_at": record.created_at.isoformat()
        } for record in records]
   

    def generate_topic_for_conversation(self, conversation_id: str) -> str:
        records = ConversationRecord.query.filter_by(
            conversation_id=conversation_id
        ).order_by(ConversationRecord.created_at.asc()).all()

        messages = [{
            "role": record.role,
            "content": record.content,
            # "created_at": record.created_at
        } for record in records]
        topic = self.llm_service.generate_topic_with_llm(messages)
        conversation = Conversation.query.get(conversation_id)
        if conversation:
            conversation.topic = topic
            db.session.commit()
            logger.info(f"为会话 {conversation_id} 生成主题: {topic}")
        else:
            logger.warning(f"会话 {conversation_id} 不存在")



    def get_topic(self,conversation_id:str)->str:
        try:
            # 查询conversations表中id等于conversation_id的记录
            conversation = Conversation.query.filter_by(id=conversation_id).first()
            if conversation:
                return conversation.topic
            else:
                return None
        except Exception as e:
            print(f"查询时出错: {e}")
            return None
